from flask import Flask, request, jsonify
import subprocess

app = Flask(__name__)

@app.route('/generate', methods=['POST'])
def generate():
    print("in generate")
    number_cha = 512
    prompt = request.json['prompt']
    command = [
        "./llama-cli",
        "-m", "/mnt/hgfs/cloudapp/qwen2-1_5b-instruct-q3_k_m.gguf",
        "-p", prompt, 
        "-n", str(number_cha),
    ]
    print(command)
    result = subprocess.run(command, capture_output=True, text=True)
    print("finish result")
    return jsonify({"response": result.stdout})

@app.route('/retrieval', methods=['POST'])
def retrieval():
    global process
    # 打印进入函数的信息
    print("in retrieval")

    # 获取请求中的 JSON 数据
    data = request.json

    # 定义命令行参数
    model_path = "/mnt/hgfs/cloudapp/qwen2-0_5b-instruct-q2_k.gguf"
    top_k = data.get('top_k', 3)  # 默认值为3
    context_file = data.get('context_file', '/home/user1/llama.cpp/llm_api_comp/graduate_employment_info.txt')  # 默认值为'content.txt'
    chunk_size = data.get('chunk_size', 100)  # 默认值为100
    chunk_separator = data.get('chunk_separator', '.')  # 默认值为'.'
    pooling = data.get('pooling', 'mean')  # 默认值为'mean'

    # 构建命令列表
    command = [
        "./llama-retrieval",
        "--model", model_path,
        "--top-k", str(top_k),
        "--context-file", context_file,
        "--chunk-size", str(chunk_size),
        "--chunk-separator", chunk_separator,
        "--pooling", pooling
    ]

    # 打印即将执行的命令
    print(command)

    # 执行命令并捕获输出
    # result = subprocess.run(command, capture_output=True, text=True)

    # 启动进程
    process = subprocess.Popen(command, stdin=subprocess.PIPE, stdout=subprocess.PIPE, text=True)

    # 初始化时读取提示信息
    initial_output = process.stdout.readline().strip()

    # 返回初始化信息
    return jsonify({"initial_output": initial_output})

@app.route('/query', methods=['POST'])
def query():
    global process

    # 打印进入函数的信息
    print("in query")

    # 获取请求中的 JSON 数据
    data = request.json

    # 获取查询词
    query = data.get('query', '')

    if not process:
        return jsonify({"error": "Retrieval process is not started. Please start it first."}), 400

    # 发送查询词
    process.stdin.write(query + '\n')
    process.stdin.flush()

    # 读取响应
    response = []
    while True:
        line = process.stdout.readline().strip()
        if line.startswith("Enter query:"):
            break
        response.append(line)

    # 返回结果
    return jsonify({"response": "\n".join(response)})

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=5000)